#include <algorithm>
#include <atomic>
#include <chrono>
#include <cstring>
#include <iostream>
#include <memory>
#include <mutex>
#include <netinet/in.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <thread>
#include <unistd.h>
#include <unordered_map>
#include <vector>

#define BUFFER_SIZE 1024
#define SELECT_TIMEOUT 3

using namespace std;
using port_t = uint16_t;
using sock_t = int;

// 前置声明
class Context;

// 数据处理函数类型定义
using select_handler = void (*)(sock_t fd, const char *buffer, ssize_t len);

// 监听器类（自动管理socket资源）
struct Listener {
	sock_t sockfd;
	port_t port;
	select_handler handler;

	Listener(sock_t s = -1, port_t p = 0) : sockfd(s), port(p), handler(nullptr) {}

	~Listener() {
		if (sockfd != -1) {
			close(sockfd);
			cout << "[监听器] 关闭socket: " << sockfd << endl;
		}
	}
};

// 网络上下文类（线程安全）
class Context {
private:
	unordered_map<sock_t, unique_ptr<Listener>> sock2listener;
	unordered_map<port_t, Listener *> port2listener;
	mutex listeners_mutex;

public:
	// 添加监听器（接管所有权）
	void addListener(unique_ptr<Listener> listener) {
		lock_guard<mutex> lock(listeners_mutex);
		auto ptr = listener.get();
		sock2listener.emplace(ptr->sockfd, move(listener));
		port2listener.emplace(ptr->port, ptr);
		cout << "[上下文] 添加监听器 - socket:" << ptr->sockfd << " 端口:" << ptr->port << endl;
	}

	// 通过socket获取监听器
	Listener *getListenerByFD(sock_t fd) {
		lock_guard<mutex> lock(listeners_mutex);
		auto it = sock2listener.find(fd);
		return it != sock2listener.end() ? it->second.get() : nullptr;
	}

	// 通过端口获取监听器
	Listener *getListenerByPort(port_t port) {
		lock_guard<mutex> lock(listeners_mutex);
		auto it = port2listener.find(port);
		return it != port2listener.end() ? it->second : nullptr;
	}

	// 移除监听器
	void removeListener(port_t port) {
		lock_guard<mutex> lock(listeners_mutex);
		auto port_it = port2listener.find(port);
		if (port_it != port2listener.end()) {
			auto *ptr = port_it->second;
			cout << "[上下文] 移除监听器 - socket:" << ptr->sockfd << " 端口:" << port << endl;
			sock2listener.erase(ptr->sockfd);
			port2listener.erase(port_it);
		}
	}

	// 获取当前所有socket描述符（线程安全副本）
	vector<sock_t> getAllSockets() {
		lock_guard<mutex> lock(listeners_mutex);
		vector<sock_t> sockets;
		for (const auto &pair: sock2listener) {
			sockets.push_back(pair.first);
		}
		return sockets;
	}
};

// UCP协议处理函数
void ucp_handler(sock_t fd, const char *buffer, ssize_t len) { cout << "[UCP处理] 来自socket " << fd << " 的数据: " << string(buffer, len) << endl; }

// 网络选择器类
class Selector {
private:
	atomic<bool> running{false};
	thread worker_thread;
	Context &context;

	void worker() {
		while (running) {
			auto sockets = context.getAllSockets(); // 获取安全副本
			if (sockets.empty()) {
				this_thread::sleep_for(chrono::seconds(1));
				continue;
			}

			fd_set readfds;
			FD_ZERO(&readfds);
			int max_fd = 0;

			for (auto sock: sockets) {
				FD_SET(sock, &readfds);
				max_fd = max(max_fd, sock);
			}

			timeval timeout{.tv_sec = SELECT_TIMEOUT};
			int ret = select(max_fd + 1, &readfds, nullptr, nullptr, &timeout);

			if (ret < 0) {
				cerr << "[选择器] select错误: " << strerror(errno) << endl;
				continue;
			}

			for (auto sock: sockets) {
				if (FD_ISSET(sock, &readfds)) {
					char buffer[BUFFER_SIZE];
					sockaddr_in client_addr{};
					socklen_t addr_len = sizeof(client_addr);

					ssize_t len = recvfrom(sock, buffer, BUFFER_SIZE, 0, (sockaddr *) &client_addr, &addr_len);

					if (len > 0) {
						if (auto *listener = context.getListenerByFD(sock)) {
							listener->handler(sock, buffer, len);
						}
					} else if (len < 0 && (errno == EBADF || errno == ENOTSOCK)) {
						// 自动清理失效socket
						if (auto *listener = context.getListenerByFD(sock)) {
							context.removeListener(listener->port);
						}
					}
				}
			}
		}
	}

public:
	Selector(Context &ctx) : context(ctx) {
		running       = true;
		worker_thread = thread(&Selector::worker, this);
	}

	~Selector() {
		running = false;
		if (worker_thread.joinable()) {
			worker_thread.join();
		}
	}
};

// 创建监听器工厂函数
unique_ptr<Listener> createListener(port_t port = 0) {
	sock_t sockfd = socket(AF_INET, SOCK_DGRAM, 0);
	if (sockfd < 0) {
		cerr << "[监听器] 创建socket失败: " << strerror(errno) << endl;
		return nullptr;
	}

	// 设置地址重用
	int reuse = 1;
	if (setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse)) < 0) {
		cerr << "[监听器] 设置SO_REUSEADDR失败: " << strerror(errno) << endl;
	}

	sockaddr_in server_addr{};
	server_addr.sin_family      = AF_INET;
	server_addr.sin_addr.s_addr = htonl(INADDR_ANY);
	server_addr.sin_port        = htons(port);

	if (::bind(sockfd, (sockaddr *) &server_addr, sizeof(server_addr)) == -1) {
		cerr << "[监听器] 绑定端口失败: " << strerror(errno) << endl;
		close(sockfd);
		return nullptr;
	}

	// 获取实际绑定端口
	socklen_t len = sizeof(server_addr);
	getsockname(sockfd, (sockaddr *) &server_addr, &len);
	auto actual_port = ntohs(server_addr.sin_port);

	cout << "[监听器] 创建成功 - socket:" << sockfd << " 端口:" << actual_port << endl;


	// return make_unique<Listener>(sockfd, actual_port);
	// 替换成c++11
	return unique_ptr<Listener>(new Listener(sockfd, actual_port));
}

int main() {
	Context context;
	vector<port_t> managed_ports;

	// 创建两个监听器
	for (int i = 0; i < 2; ++i) {
		auto listener = createListener();
		if (listener) {
			listener->handler = ucp_handler;
			managed_ports.push_back(listener->port);
			context.addListener(move(listener));
		}
	}

	// 启动选择器
	Selector selector(context);

	// 发送测试消息
	for (size_t i = 0; i < 3; ++i) {
		this_thread::sleep_for(chrono::seconds(1));

		for (auto port: managed_ports) {
			if (auto *listener = context.getListenerByPort(port)) {
				string msg = "Test message " + to_string(i + 1);
				sockaddr_in target{};
				target.sin_family      = AF_INET;
				target.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
				target.sin_port        = htons(port);

				sendto(listener->sockfd, msg.data(), msg.size(), 0, (sockaddr *) &target, sizeof(target));
			}
		}
	}

	// 清理资源
	cout << "\n[清理] 开始释放资源..." << endl;
	for (auto port: managed_ports) {
		context.removeListener(port);
	}
	cout << "[清理] 资源释放完成" << endl;

	return 0;
}
